#include <cstdio>
#include <iostream>
// #include "myType.cuh"
#include "cuda_runtime_api.h"


struct iAS {
    int idx;
    float score;
};
#define MAX_DEPTH 16
#define INSERTION_SORT 0

#define GPU_CHECK(ans)                                                         \
    { GPUAssert((ans), __FILE__, __LINE__); }
inline void GPUAssert(cudaError_t code, const char *file, int line,
                      bool abort = true) {
    if (code != cudaSuccess) {
        fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file,
                line);
        if (abort)
            exit(code);
    }
};

__device__ void selection_sort(iAS *data, int left, int right) {
    for (int i = left; i <= right; ++i) {
        float min_val = data[i].score;
        int min_idx = i;
        int temp = 0;

        // Find the smallest value in the range [left, right].
        for (int j = i + 1; j <= right; ++j) {
            float val_j = data[j].score;

            if (val_j < min_val) {
                min_idx = j;
                min_val = val_j;
            }
        }

        // Swap the values.
        if (i != min_idx) {
            data[min_idx].score = data[i].score;
            data[i].score = min_val;
            
            temp = data[min_idx].idx;
            data[min_idx].idx = data[i].idx;
            data[i].idx = temp;
        }
    }
}

__global__ void cdp_simple_quicksort(iAS *data, int left, int right,int depth, int TOPK) {
    // 当递归的深度大于设定的MAX_DEPTH或者待排序的数组长度小于设定的阈值，直接调用简单选择排序
    if (depth >= MAX_DEPTH || right - left <= INSERTION_SORT) {
        selection_sort(data, left, right);
        return;
    }

    iAS *left_ptr = data + left;
    iAS *right_ptr = data + right;
    float pivot = data[(left + right) / 2].score;
    // partition
    while (left_ptr <= right_ptr) {
        float left_val = (*left_ptr).score;
        float right_val = (*right_ptr).score;

        int left_idx = (*left_ptr).idx;
        int right_idx = (*right_ptr).idx;

        while (left_val < pivot) { // 找到第一个比pivot大的
            left_ptr++;
            left_val = (*left_ptr).score;
            left_idx = (*left_ptr).idx;
        }

        while (right_val > pivot) { // 找到第一个比pivot小的
            right_ptr--;
            right_val = (*right_ptr).score;
            right_idx = (*right_ptr).idx;
        }

        // do swap
        if (left_ptr <= right_ptr) {
            (*left_ptr).score = right_val; 
            (*left_ptr).idx = right_idx; 
            left_ptr++;
            
            (*right_ptr).score = left_val; 
            (*right_ptr).idx = left_idx; 
            right_ptr--;
        }
    }

    // recursive
    int n_right = right_ptr - data;
    int n_left = left_ptr - data;
    // Launch a new block to sort the the left part.
    
    if (left < (right_ptr - data)) {
        cudaStream_t l_stream;
        // 设置非阻塞流
        cudaStreamCreateWithFlags(&l_stream, cudaStreamNonBlocking);
        cdp_simple_quicksort<<<1, 1, 0, l_stream>>>(data, left, n_right,
                                                    depth + 1, TOPK);
        cudaStreamDestroy(l_stream);
    } 
    
    if ((left_ptr - data) < right && n_left < TOPK) {
        // Launch a new block to sort the the right part.
        cudaStream_t r_stream;
        // 设置非阻塞流
        cudaStreamCreateWithFlags(&r_stream, cudaStreamNonBlocking);
        cdp_simple_quicksort<<<1, 1, 0, r_stream>>>(data, n_left, right,
                                                    depth + 1, TOPK);
        cudaStreamDestroy(r_stream);    
    }
}

// Call the quicksort kernel from the host.
void run_qsort(iAS *data, unsigned int nitems, int TOPK) {
    // Prepare CDP for the max depth 'MAX_DEPTH'.
    GPU_CHECK(cudaDeviceSetLimit(cudaLimitDevRuntimeSyncDepth, MAX_DEPTH));

    int left = 0;
    int right = nitems - 1;
    cdp_simple_quicksort<<<1, 1>>>(data, left, right, 0, TOPK);
    GPU_CHECK(cudaDeviceSynchronize());
}

// Initialize data on the host.
void initialize_data(iAS *dst, unsigned int nitems) {
    // Fixed seed for illustration
    srand(2047);

    // Fill dst with random values
    for (unsigned i = 0; i < nitems; i++) {
        dst[i].idx = i;
        dst[i].score = rand() % (nitems + 1000);
    }
}

// Verify the results.
void print_results(int n, iAS *results_d) {
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n * sizeof(iAS), cudaMemcpyDeviceToHost));
    std::cout << "Sort data : \n";
    for (int i = 0; i < n; ++i){
        std::cout << results_h[i].idx << " " << results_h[i].score << " " << "\n";
        // if (results_h[i - 1] > results_h[i]) {
        //     std::cout << "Invalid item[" << i - 1 << "]: " << results_h[i - 1]
        //               << " greater than " << results_h[i] << std::endl;
        //     exit(EXIT_FAILURE);
        // }
    }
    std::cout << std::endl;
    delete[] results_h;
}

void check_results(int n, iAS *results_d)
{
    iAS *results_h = new iAS[n];
    GPU_CHECK(cudaMemcpy(results_h, results_d, n*sizeof(iAS), cudaMemcpyDeviceToHost));

    for (int i = 0 ; i < n ; ++i)
        if (results_h[i-1].score > results_h[i].score)
        {
            std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].idx << " greater than " << results_h[i].idx << std::endl;
            std::cout << "Invalid item[" << i-1 << "]: " << results_h[i-1].score << " greater than " << results_h[i].score << std::endl;
            exit(EXIT_FAILURE);
        }

    std::cout << "OK" << std::endl;
    delete[] results_h;
}

// Main entry point.
int main(int argc, char **argv) {
    int num_items = 20;
    bool verbose = true;
    int TOPK = 7;

    // Find/set device and get device properties
    int device = 0;
    cudaDeviceProp deviceProp;
    GPU_CHECK(cudaGetDeviceProperties(&deviceProp, device));

    if (!(deviceProp.major > 3 ||
          (deviceProp.major == 3 && deviceProp.minor >= 5))) {
        printf("GPU %d - %s  does not support CUDA Dynamic Parallelism\n Exiting.",
            device, deviceProp.name);
        return 0;
    }

    // Create input data
    iAS *h_data = 0;
    iAS *d_data = 0;

    // Allocate CPU memory and initialize data.
    std::cout << "Initializing data:" << std::endl;
    h_data = (iAS *)malloc(num_items * sizeof(iAS));
    initialize_data(h_data, num_items);
    if (verbose) {
       std::cout << "Raw  data : \n";
        for (int i = 0; i < num_items; i++)
            std::cout << h_data[i].idx << " " << h_data[i].score << "\n";
    }


    // Allocate GPU memory.
    GPU_CHECK(cudaMalloc((void **)&d_data, num_items * sizeof(iAS)));
    GPU_CHECK(cudaMemcpy(d_data, h_data, num_items * sizeof(iAS),cudaMemcpyHostToDevice));

    // Execute
    std::cout << "Running quicksort on " << num_items << " elements" << std::endl;
    run_qsort(d_data, num_items, TOPK);

    // print result
    print_results(num_items, d_data);
    // check result
    std::cout << "Validating results: ";
    // check_results(num_items, d_data);
    free(h_data);
    // GPU_CHECK(cudaFree(d_data));

    return 0;
}

/* 
编译：nvcc -o quicksort_cuda --gpu-architecture=sm_70 -rdc=truequicksort_cuda.cu
*/